{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from art.dev.get_model_config import get_model_config\n",
    "from art.local.state import ModelState\n",
    "\n",
    "config = get_model_config(\n",
    "    \"Qwen/Qwen2.5-72B-Instruct\",\n",
    "    output_dir=\"./.art/models/test\",\n",
    "    config={\n",
    "        \"engine_args\": {\n",
    "            \"enable_sleep_mode\": True,\n",
    "            \"enforce_eager\": True,\n",
    "            \"gpu_memory_utilization\": 0.9,\n",
    "        },\n",
    "        \"peft_args\": {\n",
    "            # \"use_gradient_checkpointing\": False,\n",
    "        },\n",
    "    },\n",
    ")\n",
    "state = ModelState(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "from art.local.vllm import set_vllm_log_file\n",
    "import random\n",
    "import torch\n",
    "import vllm\n",
    "\n",
    "\n",
    "set_vllm_log_file(\"./vllm.log\")\n",
    "\n",
    "num_tokens = 16384\n",
    "\n",
    "\n",
    "async def warmup(request_id: str) -> None:\n",
    "    max_tokens = random.randint(0, num_tokens * 2)\n",
    "    async for _ in state.vllm.async_engine.generate(\n",
    "        prompt={\n",
    "            \"prompt_token_ids\": torch.randint(\n",
    "                0,\n",
    "                int(state.tokenizer.vocab_size),  # type: ignore\n",
    "                (\n",
    "                    max(\n",
    "                        (\n",
    "                            (\n",
    "                                state.vllm.async_engine.engine.cache_config.num_gpu_blocks  # type: ignore\n",
    "                                * state.vllm.async_engine.engine.cache_config.block_size\n",
    "                            )\n",
    "                            // state.vllm.async_engine.engine.scheduler_config.max_num_seqs\n",
    "                        )\n",
    "                        - 16\n",
    "                        - max_tokens,\n",
    "                        1,\n",
    "                    ),\n",
    "                ),\n",
    "            ).tolist(),\n",
    "        },\n",
    "        sampling_params=vllm.SamplingParams(max_tokens=max_tokens),\n",
    "        request_id=request_id,\n",
    "    ):\n",
    "        pass\n",
    "\n",
    "\n",
    "warmup_future = asyncio.gather(\n",
    "    *(\n",
    "        warmup(f\"{i}\")\n",
    "        for i in range(state.vllm.async_engine.engine.scheduler_config.max_num_seqs)\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "warmup_future.cancel()\n",
    "await asyncio.sleep(0.01)\n",
    "try:\n",
    "    warmup_future.result()\n",
    "except asyncio.CancelledError:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "from art.local.train import train\n",
    "\n",
    "results_queue = asyncio.Queue()\n",
    "train_task = asyncio.create_task(train(state.trainer, results_queue))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from art.types import TrainConfig\n",
    "from art.local.service import TrainInputs\n",
    "import time\n",
    "from typing import cast\n",
    "\n",
    "seq_len = 32768\n",
    "batch_size = 1\n",
    "shape = (batch_size, seq_len)\n",
    "num_steps = 1\n",
    "async with state.vllm.train_mode():\n",
    "    start_time = time.time()\n",
    "    for _ in range(num_steps):\n",
    "        state.inputs_queue.put_nowait(\n",
    "            TrainInputs(\n",
    "                tokens=torch.randint(0, cast(int, state.tokenizer.vocab_size), shape),\n",
    "                group_ids=torch.randint(0, 10, shape),\n",
    "                parent_ids=torch.randint(0, 10, shape),\n",
    "                input_pos=torch.tensor([list(range(shape[1]))]),\n",
    "                assistant_mask=torch.ones(shape, dtype=torch.bool),\n",
    "                logprobs=torch.zeros(shape),\n",
    "                advantages=torch.zeros(shape),\n",
    "                weights=torch.ones(shape),\n",
    "                config=TrainConfig(lr=1e-7, kl_coef=0.01),\n",
    "            )\n",
    "        )\n",
    "        done, _ = await asyncio.wait(\n",
    "            [asyncio.create_task(results_queue.get()), train_task],\n",
    "            return_when=asyncio.FIRST_COMPLETED,\n",
    "        )\n",
    "        for task in done:\n",
    "            result = task.result()\n",
    "            # If `result` is `None`, the training task finished somehow.\n",
    "            assert result is not None, \"The training task should never finish.\"\n",
    "            results_queue.task_done()\n",
    "            display(result)\n",
    "    total_tokens = num_steps * batch_size * seq_len\n",
    "    elapsed_time = time.time() - start_time\n",
    "    tokens_per_second = total_tokens / elapsed_time\n",
    "    print(f\"Tokens per second: {tokens_per_second:.2f} tokens/s\")\n",
    "    print(f\"Total time: {elapsed_time:.2f} seconds\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
